import json
import os
import shutil
from datetime import datetime
from pathlib import Path
import argparse

class AnnotationBatchFixer:
    def __init__(self, folder_path, backup_folder=None, output_folder=None, dry_run=False):
        self.folder_path = Path(folder_path)
        self.backup_folder = Path(backup_folder) if backup_folder else self.folder_path / "backup"
        self.output_folder = Path(output_folder) if output_folder else self.folder_path / "fixed_annotations"
        self.dry_run = dry_run
        self.stats = {
            'total_files': 0,
            'json_files': 0,
            'correct_files': 0,
            'incorrect_files': 0,
            'fixed_files': 0,
            'error_files': 0,
            'copied_images': 0,
            'missing_images': 0
        }
    
    def load_annotation_file(self, file_path):
        """加载标注文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"❌ 无法加载文件 {file_path}: {e}")
            return None
    
    def save_annotation_file(self, data, file_path):
        """保存标注文件"""
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"❌ 无法保存文件 {file_path}: {e}")
            return False
    
    def get_corner_position(self, point, all_points):
        """确定点在矩形中的位置"""
        x_coords = [p[0] for p in all_points]
        y_coords = [p[1] for p in all_points]
        
        min_x, max_x = min(x_coords), max(x_coords)
        min_y, max_y = min(y_coords), max(y_coords)
        
        x, y = point
        if abs(x - min_x) < abs(x - max_x):  # 更接近左边
            if abs(y - min_y) < abs(y - max_y):  # 更接近上边
                return "左上角"
            else:  # 更接近下边
                return "左下角"
        else:  # 更接近右边
            if abs(y - min_y) < abs(y - max_y):  # 更接近上边
                return "右上角"
            else:  # 更接近下边
                return "右下角"
    
    def check_bbox_format(self, annotation_data):
        """检查bbox格式是否正确
        返回: (is_correct, error_type, details)
        """
        if not annotation_data or 'cells' not in annotation_data:
            return False, "invalid_structure", "文件结构无效"
        
        correct_count = 0
        total_count = 0
        error_details = []
        
        for cell in annotation_data['cells']:
            if not cell.get('bbox'):
                continue
                
            bbox = cell['bbox']
            if not all(key in bbox for key in ['p1', 'p2', 'p3', 'p4']):
                continue
                
            p1, p2, p3, p4 = bbox['p1'], bbox['p2'], bbox['p3'], bbox['p4']
            all_points = [p1, p2, p3, p4]
            
            # 检查每个点是否在正确位置
            p1_pos = self.get_corner_position(p1, all_points)  # 应该是左上角
            p2_pos = self.get_corner_position(p2, all_points)  # 应该是右上角
            p3_pos = self.get_corner_position(p3, all_points)  # 应该是右下角
            p4_pos = self.get_corner_position(p4, all_points)  # 应该是左下角
            
            is_correct = (p1_pos == "左上角" and p2_pos == "右上角" and 
                         p3_pos == "右下角" and p4_pos == "左下角")
            
            if is_correct:
                correct_count += 1
            else:
                error_details.append({
                    'cell_ind': cell.get('cell_ind', total_count),
                    'expected': "p1(左上) p2(右上) p3(右下) p4(左下)",
                    'actual': f"p1({p1_pos}) p2({p2_pos}) p3({p3_pos}) p4({p4_pos})"
                })
            
            total_count += 1
        
        if total_count == 0:
            return False, "no_bbox", "文件中没有有效的bbox数据"
        
        success_rate = correct_count / total_count
        if success_rate >= 0.9:  # 90%以上正确认为是正确文件
            return True, "correct", f"正确率: {success_rate:.1%}"
        else:
            return False, "incorrect_order", f"正确率: {success_rate:.1%}, 错误详情: {error_details[:3]}"
    
    def fix_bbox_order(self, annotation_data):
        """修复bbox点顺序"""
        if not annotation_data or 'cells' not in annotation_data:
            return None
        
        fixed_cells = 0
        for cell in annotation_data['cells']:
            if not cell.get('bbox'):
                continue
                
            bbox = cell['bbox']
            if not all(key in bbox for key in ['p1', 'p2', 'p3', 'p4']):
                continue
            
            # 获取原始点
            p1_orig = bbox['p1']  # 可能是右上角
            p2_orig = bbox['p2']  # 可能是左上角  
            p3_orig = bbox['p3']  # 可能是左下角
            p4_orig = bbox['p4']  # 可能是右下角
            
            all_points = [p1_orig, p2_orig, p3_orig, p4_orig]
            
            # 找到每个角的正确点
            left_top = None
            right_top = None
            right_bottom = None
            left_bottom = None
            
            for point in all_points:
                pos = self.get_corner_position(point, all_points)
                if pos == "左上角":
                    left_top = point
                elif pos == "右上角":
                    right_top = point
                elif pos == "右下角":
                    right_bottom = point
                elif pos == "左下角":
                    left_bottom = point
            
            # 重新分配点位置为标准的逆时针顺序
            if all(p is not None for p in [left_top, right_top, right_bottom, left_bottom]):
                bbox['p1'] = left_top      # 左上角
                bbox['p2'] = right_top     # 右上角
                bbox['p3'] = right_bottom  # 右下角
                bbox['p4'] = left_bottom   # 左下角
                fixed_cells += 1
        
        # 更新质量状态
        if 'quality' in annotation_data:
            annotation_data['quality'] = "已修复"
        
        return annotation_data, fixed_cells
    
    def get_image_path_from_json(self, json_file_path):
        """根据JSON文件名推断对应的图片文件路径"""
        # 标注文件命名格式：图片名_table_annotation.json
        json_name = json_file_path.stem
        if json_name.endswith('_table_annotation'):
            # 提取图片名（去掉_table_annotation后缀）
            image_name = json_name[:-17]  # 去掉'_table_annotation'

            # 常见图片格式
            image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']

            for ext in image_extensions:
                image_path = json_file_path.parent / f"{image_name}{ext}"
                if image_path.exists():
                    return image_path

        return None

    def copy_files_to_output(self, json_file_path, image_file_path=None):
        """将JSON文件和对应的图片复制到输出文件夹"""
        if self.dry_run:
            return True

        try:
            # 确保输出文件夹存在
            self.output_folder.mkdir(parents=True, exist_ok=True)

            # 复制JSON文件
            output_json_path = self.output_folder / json_file_path.name
            shutil.copy2(json_file_path, output_json_path)

            # 复制图片文件（如果存在）
            if image_file_path and image_file_path.exists():
                output_image_path = self.output_folder / image_file_path.name
                shutil.copy2(image_file_path, output_image_path)
                self.stats['copied_images'] += 1
                return True, f"已复制JSON和图片文件到 {self.output_folder}"
            else:
                self.stats['missing_images'] += 1
                return True, f"已复制JSON文件到 {self.output_folder}，但未找到对应图片"

        except Exception as e:
            print(f"❌ 复制文件失败: {e}")
            return False, f"复制失败: {e}"

    def create_backup(self, file_path):
        """创建备份文件"""
        if self.dry_run:
            return True

        try:
            # 确保备份文件夹存在
            self.backup_folder.mkdir(parents=True, exist_ok=True)

            # 创建备份文件名
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            backup_name = f"{file_path.stem}_backup_{timestamp}{file_path.suffix}"
            backup_path = self.backup_folder / backup_name

            shutil.copy2(file_path, backup_path)
            return True
        except Exception as e:
            print(f"❌ 创建备份失败 {file_path}: {e}")
            return False
    
    def process_file(self, file_path):
        """处理单个文件"""
        print(f"\n📄 处理文件: {file_path.name}")

        # 查找对应的图片文件
        image_path = self.get_image_path_from_json(file_path)
        if image_path:
            print(f"🖼️  找到对应图片: {image_path.name}")
        else:
            print(f"⚠️  未找到对应图片文件")

        # 加载文件
        annotation_data = self.load_annotation_file(file_path)
        if annotation_data is None:
            self.stats['error_files'] += 1
            return False

        # 检查格式
        is_correct, error_type, details = self.check_bbox_format(annotation_data)

        # 处理正确的文件
        if is_correct:
            print(f"✅ 文件格式正确: {details}")
            self.stats['correct_files'] += 1

            # 复制到输出文件夹
            if self.dry_run:
                print("🔍 [预览模式] 将会复制此文件到输出文件夹")
            else:
                success, message = self.copy_files_to_output(file_path, image_path)
                if success:
                    print(f"📁 {message}")
                else:
                    print(f"❌ {message}")
                    return False
            return True

        # 处理错误的文件
        print(f"⚠️  发现问题: {error_type} - {details}")
        self.stats['incorrect_files'] += 1

        if error_type in ['invalid_structure', 'no_bbox']:
            print(f"❌ 无法修复: {error_type}")
            self.stats['error_files'] += 1
            return False

        # 尝试修复
        if self.dry_run:
            print("🔍 [预览模式] 将会修复此文件并复制到输出文件夹")
            self.stats['fixed_files'] += 1
            return True

        # 创建备份
        if not self.create_backup(file_path):
            return False

        # 执行修复
        try:
            fixed_data, fixed_cells = self.fix_bbox_order(annotation_data)
            if fixed_data is None:
                print("❌ 修复失败")
                return False

            # 保存修复后的文件到输出文件夹
            output_json_path = self.output_folder / file_path.name
            self.output_folder.mkdir(parents=True, exist_ok=True)

            if self.save_annotation_file(fixed_data, output_json_path):
                print(f"✅ 修复成功: 修复了 {fixed_cells} 个单元格")

                # 复制图片文件到输出文件夹
                if image_path and image_path.exists():
                    output_image_path = self.output_folder / image_path.name
                    shutil.copy2(image_path, output_image_path)
                    self.stats['copied_images'] += 1
                    print(f"📁 已将修复后的JSON和图片文件保存到 {self.output_folder}")
                else:
                    self.stats['missing_images'] += 1
                    print(f"📁 已将修复后的JSON文件保存到 {self.output_folder}，但未找到对应图片")

                self.stats['fixed_files'] += 1
                return True
            else:
                print("❌ 保存修复文件失败")
                return False

        except Exception as e:
            print(f"❌ 修复过程出错: {e}")
            self.stats['error_files'] += 1
            return False
    
    def run(self):
        """批量处理文件夹中的所有JSON文件"""
        print("="*80)
        print("表格标注文件批量修复工具")
        print("="*80)
        print(f"📁 处理文件夹: {self.folder_path}")
        print(f"💾 备份文件夹: {self.backup_folder}")
        print(f"� 输出文件夹: {self.output_folder}")
        print(f"�🔍 预览模式: {'是' if self.dry_run else '否'}")
        print("-"*80)
        
        # 查找所有JSON文件
        json_files = list(self.folder_path.glob("*.json"))
        self.stats['total_files'] = len(list(self.folder_path.iterdir()))
        self.stats['json_files'] = len(json_files)
        
        if not json_files:
            print("❌ 文件夹中没有找到JSON文件")
            return
        
        print(f"📊 找到 {len(json_files)} 个JSON文件")
        
        # 处理每个文件
        for file_path in json_files:
            self.process_file(file_path)
        
        # 输出统计结果
        self.print_summary()
    
    def print_summary(self):
        """打印处理结果摘要"""
        print("\n" + "="*80)
        print("处理结果摘要")
        print("="*80)
        print(f"📁 总文件数: {self.stats['total_files']}")
        print(f"📄 JSON文件数: {self.stats['json_files']}")
        print(f"✅ 格式正确: {self.stats['correct_files']}")
        print(f"⚠️  格式错误: {self.stats['incorrect_files']}")
        print(f"🔧 已修复: {self.stats['fixed_files']}")
        print(f"❌ 处理失败: {self.stats['error_files']}")
        print(f"🖼️  已复制图片: {self.stats['copied_images']}")
        print(f"📷 缺失图片: {self.stats['missing_images']}")

        if self.stats['json_files'] > 0:
            correct_rate = (self.stats['correct_files'] / self.stats['json_files']) * 100
            fix_rate = (self.stats['fixed_files'] / max(self.stats['incorrect_files'], 1)) * 100
            image_rate = (self.stats['copied_images'] / self.stats['json_files']) * 100
            print(f"📈 原始正确率: {correct_rate:.1f}%")
            print(f"🔧 修复成功率: {fix_rate:.1f}%")
            print(f"🖼️  图片匹配率: {image_rate:.1f}%")

        if not self.dry_run:
            if self.stats['fixed_files'] > 0 or self.stats['correct_files'] > 0:
                print(f"\n📂 输出文件夹: {self.output_folder}")
            if self.stats['fixed_files'] > 0:
                print(f"💾 备份文件保存在: {self.backup_folder}")

        print("="*80)

def main():
    parser = argparse.ArgumentParser(description='批量修复表格标注文件的bbox点顺序')
    parser.add_argument('folder',default=r"F:\workspace\datasets\Relabel_TALOCRTable\borderless_tables", help='包含标注文件的文件夹路径')
    parser.add_argument('--backup', help='备份文件夹路径（默认为输入文件夹下的backup子文件夹）')
    parser.add_argument('--output', help='输出文件夹路径（默认为输入文件夹下的fixed_annotations子文件夹）')
    parser.add_argument('--dry-run', action='store_true', help='预览模式，不实际修改文件')

    args = parser.parse_args()

    # 检查文件夹是否存在
    if not os.path.exists(args.folder):
        print(f"❌ 文件夹不存在: {args.folder}")
        return

    # 创建修复器并运行
    fixer = AnnotationBatchFixer(
        folder_path=args.folder,
        backup_folder=args.backup,
        output_folder=args.output,
        dry_run=args.dry_run
    )

    fixer.run()

if __name__ == "__main__":
    # 如果没有命令行参数，使用当前目录
    import sys
    if len(sys.argv) == 1:
        print("使用当前目录进行处理...")
        fixer = AnnotationBatchFixer(
            folder_path=".",
            dry_run=True  # 默认预览模式
        )
        fixer.run()
    else:
        main()
